import json
import heapq
import matplotlib.pyplot as plt
import os
import base64
import io
import datetime
plt.rcParams["font.family"] = "sans-serif"
plt.rcParams["font.sans-serif"] = ["Microsoft YaHei", "SimHei", "SimSun", "Arial Unicode MS", "sans-serif"]
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
import networkx as nx
from mcp.server.fastmcp import FastMCP

# 读取图结构（JSON 格式）
script_dir = os.path.dirname(os.path.abspath(__file__))
graph_path = os.path.join(script_dir, 'graph.json')
config_path = os.path.join(script_dir, 'config.json')

# 读取配置文件
with open(config_path, 'r', encoding='utf-8') as f:
    config_data = json.load(f)

with open(graph_path, 'r', encoding='utf-8') as f:
    graph_data = json.load(f)

# 构建图
G = nx.Graph()
for node, edges in graph_data.items():
    for neighbor, weight in edges.items():
        G.add_edge(node, neighbor, weight=float(weight))

# Floyd 全图最短路径预计算
floyd_predecessors, floyd_dist = nx.floyd_warshall_predecessor_and_distance(G, weight='weight')
floyd_path = floyd_predecessors

def dijkstra_path(G, start, end):
    queue, visited, dist, prev = [(0, start)], set(), {start: 0}, {}
    while queue:
        cost, node = heapq.heappop(queue)
        if node in visited:
            continue
        visited.add(node)
        for neighbor in G.neighbors(node):
            weight = G[node][neighbor]['weight']
            new_cost = cost + weight
            if neighbor not in dist or new_cost < dist[neighbor]:
                dist[neighbor] = new_cost
                prev[neighbor] = node
                heapq.heappush(queue, (new_cost, neighbor))
    # 重构路径
    path = []
    if end in dist:
        current = end
        while current:
            path.insert(0, current)
            current = prev.get(current)
    return path, dist.get(end, float('inf')) if path else ([], float('inf'))

def floyd_reconstruct_path(start, end):
    if start not in floyd_dist or end not in floyd_dist[start]:
        return [], float('inf')
    # 使用networkx的最短路径函数
    path = nx.reconstruct_path(start, end, floyd_path)
    return path, floyd_dist[start][end]

def query_path(start, end, algorithm=None):
    if algorithm == 'floyd':
        path, dist = floyd_reconstruct_path(start, end)
        used_algo = 'Floyd'
    elif algorithm == 'dijkstra':
        path, dist = dijkstra_path(G, start, end)
        used_algo = 'Dijkstra'
    else:
        # 智能判断是否已计算 Floyd，否则使用 Dijkstra
        if start in floyd_dist and end in floyd_dist[start]:
            path, dist = floyd_reconstruct_path(start, end)
            used_algo = 'Floyd (auto)'
        else:
            path, dist = dijkstra_path(G, start, end)
            used_algo = 'Dijkstra (auto)'
    return path, dist, used_algo

def visualize_path_to_file(path, start, end, algorithm, distance):
    """生成路径可视化图并保存为文件，返回文件路径"""
    # 从配置文件读取可视化文件夹路径
    vis_folder = config_data.get('visualization_folder', '可视化图')
    vis_dir = os.path.join(script_dir, vis_folder)
    if os.path.exists(vis_dir):
        # 可视化图文件夹已存在，直接使用
        pass
    else:
        # 可视化图文件夹不存在，创建新目录
        os.makedirs(vis_dir)
    
    # 生成唯一文件名
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"route_{start}_{end}_{timestamp}.png"
    filepath = os.path.join(vis_dir, filename)
    
    # 创建图形
    plt.figure(figsize=(12, 8))
    pos = nx.spring_layout(G, seed=42)  # 固定随机种子保证布局一致
    
    # 绘制所有节点和边
    nx.draw(G, pos, with_labels=True, node_color='lightblue', 
            node_size=800, font_size=10, font_weight='bold')
    
    # 高亮显示最短路径
    if path and len(path) > 1:
        edges = list(zip(path, path[1:]))
        nx.draw_networkx_edges(G, pos, edgelist=edges, edge_color='red', width=3)
        nx.draw_networkx_nodes(G, pos, nodelist=path, node_color='orange', node_size=1000)
        
        # 高亮起点和终点
        nx.draw_networkx_nodes(G, pos, nodelist=[start], node_color='green', node_size=1200)
        nx.draw_networkx_nodes(G, pos, nodelist=[end], node_color='red', node_size=1200)
    
    # 添加标题和信息
    title = f"最优路径: {start} → {end}\n算法: {algorithm} | 总距离: {distance:.1f}km"
    if path:
        title += f"\n路径: {' → '.join(path)}"
    plt.title(title, fontsize=14, pad=20)
    
    # 保存图片
    plt.tight_layout()
    plt.savefig(filepath, dpi=300, bbox_inches='tight')
    plt.close()  # 关闭图形以释放内存
    
    return filepath

def visualize_path_to_base64(path, start, end, algorithm, distance):
    """生成路径可视化图并返回base64编码"""
    # 创建图形
    plt.figure(figsize=(12, 8))
    pos = nx.spring_layout(G, seed=42)  # 固定随机种子保证布局一致
    
    # 绘制所有节点和边
    nx.draw(G, pos, with_labels=True, node_color='lightblue', 
            node_size=800, font_size=10, font_weight='bold')
    
    # 高亮显示最短路径
    if path and len(path) > 1:
        edges = list(zip(path, path[1:]))
        nx.draw_networkx_edges(G, pos, edgelist=edges, edge_color='red', width=3)
        nx.draw_networkx_nodes(G, pos, nodelist=path, node_color='orange', node_size=1000)
        
        # 高亮起点和终点
        nx.draw_networkx_nodes(G, pos, nodelist=[start], node_color='green', node_size=1200)
        nx.draw_networkx_nodes(G, pos, nodelist=[end], node_color='red', node_size=1200)
    
    # 添加标题和信息
    title = f"最优路径: {start} → {end}\n算法: {algorithm} | 总距离: {distance:.1f}km"
    if path:
        title += f"\n路径: {' → '.join(path)}"
    plt.title(title, fontsize=14, pad=20)
    
    # 转换为base64
    buffer = io.BytesIO()
    plt.tight_layout()
    plt.savefig(buffer, format='png', dpi=300, bbox_inches='tight')
    buffer.seek(0)
    image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
    plt.close()  # 关闭图形以释放内存
    
    return image_base64

def visualize_path(path):
    """原有的可视化函数，用于直接显示"""
    pos = nx.spring_layout(G)
    plt.figure(figsize=(10, 6))
    nx.draw(G, pos, with_labels=True, node_color='lightgray')
    if path:
        edges = list(zip(path, path[1:]))
        nx.draw_networkx_edges(G, pos, edgelist=edges, edge_color='red', width=2)
        nx.draw_networkx_nodes(G, pos, nodelist=path, node_color='orange')
    plt.title("最短路径可视化")
    plt.show()

# 创建MCP服务器实例
mcp = FastMCP("Route Optimization Server")

@mcp.tool()
def query_route(start: str, end: str, algorithm: str = None, visualization: str = "base64") -> dict:
    """Query the optimal route between two cities with visualization
    
    Args:
        start: 起点城市名称
        end: 终点城市名称
        algorithm: 算法选择 (floyd/dijkstra/auto)
        visualization: 可视化输出格式 (base64/file/none)
    
    Returns:
        包含路径信息和可视化结果的字典
    """
    path, distance, used_algo = query_path(start, end, algorithm)
    
    result = {
        "algorithm": used_algo,
        "path": path,
        "distance": distance,
        "route_summary": f"{start} → {end} ({distance:.1f}km)"
    }
    
    # 根据visualization参数生成可视化
    if visualization == "base64":
        image_base64 = visualize_path_to_base64(path, start, end, used_algo, distance)
        result["visualization"] = {
            "type": "base64",
            "data": image_base64,
            "format": "png"
        }
    elif visualization == "file":
        filepath = visualize_path_to_file(path, start, end, used_algo, distance)
        result["visualization"] = {
            "type": "file",
            "path": filepath,
            "format": "png"
        }
    elif visualization == "none":
        result["visualization"] = None
    
    return result

@mcp.tool()
def analyze_route_comprehensive(query: str) -> dict:
    """Comprehensive route analysis with detailed results and visualization
    
    Args:
        query: 查询字符串，支持格式如"起点：北京，终点：重庆"或"北京到重庆"
    
    Returns:
        包含详细分析结果、性能对比和可视化的完整报告
    """
    import re
    import time
    
    # 解析查询字符串
    start, end = None, None
    
    # 匹配"起点：xxx，终点：xxx"格式
    pattern1 = r'起点[：:](\S+)[，,]\s*终点[：:](\S+)'
    match1 = re.search(pattern1, query)
    if match1:
        start, end = match1.groups()
    else:
        # 匹配"xxx到xxx"格式
        pattern2 = r'(\S+)到(\S+)'
        match2 = re.search(pattern2, query)
        if match2:
            start, end = match2.groups()
        else:
            # 匹配"xxx-xxx"或"xxx→xxx"格式
            pattern3 = r'(\S+)[-→](\S+)'
            match3 = re.search(pattern3, query)
            if match3:
                start, end = match3.groups()
    
    if not start or not end:
        return {
            "error": "无法解析查询格式",
            "supported_formats": [
                "起点：北京，终点：重庆",
                "北京到重庆",
                "北京→重庆",
                "北京-重庆"
            ]
        }
    
    # 去除可能的空格和标点
    start = start.strip()
    end = end.strip()
    
    # 1. 基本路径查询
    path_auto, distance_auto, algo_auto = query_path(start, end)
    path_floyd, distance_floyd, algo_floyd = query_path(start, end, 'floyd')
    path_dijkstra, distance_dijkstra, algo_dijkstra = query_path(start, end, 'dijkstra')
    
    # 2. 性能测试
    # Floyd性能测试
    start_time = time.time()
    for _ in range(100):
        query_path(start, end, 'floyd')
    floyd_time = time.time() - start_time
    
    # Dijkstra性能测试
    start_time = time.time()
    for _ in range(100):
        query_path(start, end, 'dijkstra')
    dijkstra_time = time.time() - start_time
    
    # 3. 生成可视化
    # 文件可视化
    filepath = visualize_path_to_file(path_auto, start, end, algo_auto, distance_auto)
    
    # Base64可视化
    image_base64 = visualize_path_to_base64(path_auto, start, end, algo_auto, distance_auto)
    
    # 4. 构建完整结果
    result = {
        "query_parsed": {
            "original_query": query,
            "start_city": start,
            "end_city": end
        },
        "route_analysis": {
            "auto_selection": {
                "algorithm": algo_auto,
                "path": path_auto,
                "distance": distance_auto,
                "route_summary": f"{start} → {end} ({distance_auto:.1f}km)"
            },
            "floyd_algorithm": {
                "algorithm": algo_floyd,
                "path": path_floyd,
                "distance": distance_floyd,
                "route_summary": f"{start} → {end} ({distance_floyd:.1f}km)"
            },
            "dijkstra_algorithm": {
                "algorithm": algo_dijkstra,
                "path": path_dijkstra,
                "distance": distance_dijkstra,
                "route_summary": f"{start} → {end} ({distance_dijkstra:.1f}km)"
            }
        },
        "performance_analysis": {
            "floyd_time_100_queries": f"{floyd_time:.4f}秒",
            "dijkstra_time_100_queries": f"{dijkstra_time:.4f}秒",
            "performance_comparison": f"Floyd {'更快' if floyd_time < dijkstra_time else '更慢'}",
            "speed_ratio": f"{dijkstra_time/floyd_time:.1f}x" if floyd_time > 0 else "N/A"
        },
        "visualization": {
            "file_output": {
                "type": "file",
                "path": filepath,
                "exists": os.path.exists(filepath),
                "size_bytes": os.path.getsize(filepath) if os.path.exists(filepath) else 0,
                "format": "png"
            },
            "base64_output": {
                "type": "base64",
                "data": image_base64,
                "data_length": len(image_base64),
                "format": "png",
                "preview": image_base64[:50] + "..."
            }
        },
        "features_summary": {
            "path_planning": [
                "智能算法选择（自动选择Floyd或Dijkstra）",
                "Floyd-Warshall全源最短路径算法",
                "Dijkstra单源最短路径算法",
                "返回最优路径和总距离"
            ],
            "visualization_features": [
                "生成高分辨率PNG图片文件（300 DPI）",
                "返回Base64编码的图片数据",
                "支持中文城市名称显示",
                "高亮显示最优路径（红色边，橙色节点）",
                "特殊标记起点（绿色）和终点（红色）",
                "显示算法类型、总距离和完整路径信息"
            ],
            "technical_features": [
                "MCP工具接口，支持多种可视化模式",
                "灵活的参数配置（算法选择、可视化格式）",
                "结构化的JSON返回结果",
                "Web服务接口支持",
                "高性能算法实现"
            ]
        }
    }
    
    return result



# MCP服务器和Web服务器启动
if __name__ == '__main__':
    # 读取配置文件
    with open(config_path, 'r', encoding='utf-8') as f:
        config = json.load(f)
    
    # 启动MCP服务
    import threading
    threading.Thread(target=mcp.run, kwargs={"transport": "stdio"}, daemon=True).start()
    
    # 启动Web服务器
    from http.server import HTTPServer, BaseHTTPRequestHandler
    import json

    class RequestHandler(BaseHTTPRequestHandler):
        def do_GET(self):
            if self.path.startswith('/query'):
                # 解析查询参数
                from urllib.parse import urlparse, parse_qs
                query = parse_qs(urlparse(self.path).query)
                start = query.get('start', [None])[0]
                end = query.get('end', [None])[0]
                algorithm = query.get('algorithm', [None])[0]
                
                if start and end:
                    path, distance, used = query_path(start, end, algorithm)
                    response = {
                        'algorithm': used,
                        'path': path,
                        'distance': distance
                    }
                    self.send_response(200)
                    self.send_header('Content-type', 'application/json; charset=utf-8')
                    self.end_headers()
                    self.wfile.write(json.dumps(response, ensure_ascii=False).encode('utf-8'))
                    return
                
            self.send_response(400)
            self.end_headers()
            self.wfile.write(b'Invalid request')

    server_address = (config['host'], config['port'])
    httpd = HTTPServer(server_address, RequestHandler)
    print(f"Web server running on http://{config['host']}:{config['port']}")
    httpd.serve_forever()
